{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vkdnLiKk71g-"
      },
      "source": [
        "##### Copyright 2022 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "0asMuNro71hA"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jXslvcRocA-0"
      },
      "source": [
        "# Composing Learning Algorithms"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0XBJJIqwcXKd"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/federated/tutorials/composing_learning_algorithms\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/federated/blob/v0.84.0/docs/tutorials/composing_learning_algorithms.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/federated/blob/v0.84.0/docs/tutorials/composing_learning_algorithms.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/federated/docs/tutorials/composing_learning_algorithms.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MnUwFbCAKB2r"
      },
      "source": [
        "## Before you start\n",
        "\n",
        "Before you start, please run the following to make sure that your environment is\n",
        "correctly setup. If you don't see a greeting, please refer to the\n",
        "[Installation](../install.md) guide for instructions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZrGitA_KnRO0"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install --quiet --upgrade tensorflow-federated"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "HGTM6tWOLo8M"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_federated as tff"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yr3ztf28fa1F"
      },
      "source": [
        "**NOTE**: This colab has been verified to work with the [latest released version](https://github.com/tensorflow/federated#compatibility) of the `tensorflow_federated` pip package, but the Tensorflow Federated project is still in pre-release development and may not work on `main`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CFlTaHe0jV2S"
      },
      "source": [
        "# Composing Learning Algorithms"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3zQlyijofSzI"
      },
      "source": [
        "The [Building Your Own Federated Learning Algorithm Tutorial](https://github.com/tensorflow/federated/blob/v0.84.0/docs/tutorials/building_your_own_federated_learning_algorithm.ipynb) used TFF's federated core to directly implement a version of the Federated Averaging (FedAvg) algorithm.\n",
        "\n",
        "In this tutorial, you will use federated learning components in TFF's API to build federated learning algorithms in a modular manner, without having to re-implement everything from scratch.\n",
        "\n",
        "For the purposes of this tutorial, you will implement a variant of FedAvg that employs gradient clipping through local training."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fHwcFnLAjqcG"
      },
      "source": [
        "## Learning Algorithm Building Blocks\n",
        "\n",
        "At a high level, many learning algorithms can be separated into 4 separate components, referred to as **building blocks**. These are as follows:\n",
        "\n",
        "1. Distributor (ie. server-to-client communication)\n",
        "1. Client work (ie. local client computation)\n",
        "1. Aggregator (ie. client-to-server communication)\n",
        "1. Finalizer (ie. server computation using aggregated client outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YwhOtjlvjboB"
      },
      "source": [
        "While the [Building Your Own Federated Learning Algorithm Tutorial](https://github.com/tensorflow/federated/blob/v0.84.0/docs/tutorials/building_your_own_federated_learning_algorithm.ipynb) implemented all of these building blocks from scratch, this is often unnecessary. Instead, you can re-use building blocks from similar algorithms.\n",
        "\n",
        "In this case, to implement FedAvg with gradient clipping, you only need to modify the **client work** building block. The remaining blocks can be identical to what is used in \"vanilla\" FedAvg."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LMnd0RvGlGjK"
      },
      "source": [
        "# Implementing the Client Work\n",
        "\n",
        "First, let's write TF logic that does local model training with gradient clipping. For simplicity, gradients will be clipped have norm at most 1."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-lqZ-c4MphTU"
      },
      "source": [
        "## TF Logic"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "pIw7QQCqltdV"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def client_update(\n",
        "    model: tff.learning.models.FunctionalModel,\n",
        "    dataset: tf.data.Dataset,\n",
        "    initial_weights: tff.learning.models.ModelWeights,\n",
        "    client_optimizer: tff.learning.optimizers.Optimizer,\n",
        "):\n",
        "  \"\"\"Performs training (using the initial server model weights) on the client's dataset.\"\"\"\n",
        "  # Keep track of the number of examples.\n",
        "  num_examples = 0.0\n",
        "  # Use the client_optimizer to update the local model.\n",
        "  trainable_weights, non_trainable_weights = (\n",
        "      initial_weights.trainable,\n",
        "      initial_weights.non_trainable,\n",
        "  )\n",
        "  optimizer_state = client_optimizer.initialize(\n",
        "      tf.nest.map_structure(lambda x: tf.TensorSpec, trainable_weights)\n",
        "  )\n",
        "  for batch in dataset:\n",
        "    x, y = batch\n",
        "    with tf.GradientTape() as tape:\n",
        "      tape.watch(trainable_weights)\n",
        "      logits = model.predict_on_batch(\n",
        "          model_weights=(trainable_weights, non_trainable_weights),\n",
        "          x=x,\n",
        "          training=True,\n",
        "      )\n",
        "      num_examples += tf.cast(tf.shape(y)[0], tf.float32)\n",
        "      loss = model.loss(output=logits, label=y)\n",
        "\n",
        "    # Compute the corresponding gradient\n",
        "    grads = tape.gradient(loss, trainable_weights)\n",
        "\n",
        "    # Compute the gradient norm and clip\n",
        "    gradient_norm = tf.linalg.global_norm(grads)\n",
        "    if gradient_norm \u003e 1:\n",
        "      grads = tf.nest.map_structure(lambda x: x / gradient_norm, grads)\n",
        "\n",
        "    # Apply the gradient using a client optimizer.\n",
        "    optimizer_state, trainable_weights = client_optimizer.next(\n",
        "        optimizer_state, trainable_weights, grads\n",
        "    )\n",
        "\n",
        "  # Compute the difference between the initial weights and the client weights\n",
        "  client_update = tf.nest.map_structure(\n",
        "      tf.subtract, trainable_weights, initial_weights[0]\n",
        "  )\n",
        "\n",
        "  return tff.learning.templates.ClientResult(\n",
        "      update=client_update, update_weight=num_examples\n",
        "  )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fe_emK8LpQe0"
      },
      "source": [
        "There are a few important points about the code above. First, it keeps track of the number of examples seen, as this will constitute the *weight* of the client update (when computing an average across clients).\n",
        "\n",
        "Second, it uses [`tff.learning.templates.ClientResult`](https://www.tensorflow.org/federated/api_docs/python/tff/learning/templates/ClientResult) to package the output. This return type is used to standardize client work building blocks in `tff.learning`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l5aKjB1Vpiv3"
      },
      "source": [
        "## Creating a ClientWorkProcess"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6IvXUJAzm8ab"
      },
      "source": [
        "While the TF logic above will do local training with clipping, it still needs to be wrapped in TFF code in order to create the necessary building block.\n",
        "\n",
        "Specifically, the 4 building blocks are represented as a [`tff.templates.MeasuredProcess`](https://www.tensorflow.org/federated/api_docs/python/tff/templates/MeasuredProcess). This means that all 4 blocks have both an `initialize` and `next` function used to instantiate and run the computation.\n",
        "\n",
        "This allows each building block to keep track of its own **state** (stored at the server) as needed to perform its operations. While it will not be used in this tutorial, it can be used for things like tracking how many iterations have occurred, or keeping track of optimizer states.\n",
        "\n",
        "Client work TF logic should generally be wrapped as a [`tff.learning.templates.ClientWorkProcess`](https://www.tensorflow.org/federated/api_docs/python/tff/learning/templates/ClientWorkProcess), which codifies the expected types going into and out of the client's local training. It can be parameterized by a model and optimizer, as below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "X-I-jPsZmmMy"
      },
      "outputs": [],
      "source": [
        "def build_gradient_clipping_client_work(\n",
        "    model: tff.learning.models.FunctionalModel,\n",
        "    optimizer: tff.learning.optimizers.Optimizer,\n",
        ") -\u003e tff.learning.templates.ClientWorkProcess:\n",
        "  \"\"\"Creates a client work process that uses gradient clipping.\"\"\"\n",
        "  data_type = tff.SequenceType(tff.types.tensorflow_to_type(model.input_spec))\n",
        "  model_weights_type = tff.types.to_type(\n",
        "      tf.nest.map_structure(\n",
        "          lambda arr: tff.types.TensorType(shape=arr.shape, dtype=arr.dtype),\n",
        "          tff.learning.models.ModelWeights(*model.initial_weights),\n",
        "      )\n",
        "  )\n",
        "\n",
        "  @tff.federated_computation\n",
        "  def initialize_fn():\n",
        "    return tff.federated_value((), tff.SERVER)\n",
        "\n",
        "  @tff.tensorflow.computation(model_weights_type, data_type)\n",
        "  def client_update_computation(model_weights, dataset):\n",
        "    return client_update(model, dataset, model_weights, optimizer)\n",
        "\n",
        "  @tff.federated_computation(\n",
        "      initialize_fn.type_signature.result,\n",
        "      tff.FederatedType(model_weights_type, tff.CLIENTS),\n",
        "      tff.FederatedType(data_type, tff.CLIENTS),\n",
        "  )\n",
        "  def next_fn(state, model_weights, client_dataset):\n",
        "    client_result = tff.federated_map(\n",
        "        client_update_computation, (model_weights, client_dataset)\n",
        "    )\n",
        "    # Return empty measurements, though a more complete algorithm might\n",
        "    # measure something here.\n",
        "    measurements = tff.federated_value((), tff.SERVER)\n",
        "    return tff.templates.MeasuredProcessOutput(\n",
        "        state, client_result, measurements\n",
        "    )\n",
        "\n",
        "  return tff.learning.templates.ClientWorkProcess(initialize_fn, next_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NMUX0d0Sx1Gq"
      },
      "source": [
        "# Composing a Learning Algorithm\n",
        "\n",
        "Let's put the client work above into a full-fledged algorithm. First, let's set up our data and model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hQ_N9XbULo8P"
      },
      "source": [
        "## Preparing the input data\n",
        "Load and preprocess the EMNIST dataset included in TFF. For more details, see the [image classification](federated_learning_for_image_classification.ipynb) tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "-WdnFluLLo8P"
      },
      "outputs": [],
      "source": [
        "emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kq8893GogB8E"
      },
      "source": [
        "In order to feed the dataset into our model, the data is flattened and converted into tuples of the form `(flattened_image_vector, label)`.\n",
        "\n",
        "Let's select a small number of clients, and apply the preprocessing above to their datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "Blrh8zJgLo8R"
      },
      "outputs": [],
      "source": [
        "NUM_CLIENTS = 10\n",
        "BATCH_SIZE = 20\n",
        "\n",
        "def preprocess(dataset):\n",
        "\n",
        "  def batch_format_fn(element):\n",
        "    \"\"\"Flatten a batch of EMNIST data and return a (features, label) tuple.\"\"\"\n",
        "    return (tf.reshape(element['pixels'], [-1, 784]),\n",
        "            tf.reshape(element['label'], [-1, 1]))\n",
        "\n",
        "  return dataset.batch(BATCH_SIZE).map(batch_format_fn)\n",
        "\n",
        "client_ids = sorted(emnist_train.client_ids)[:NUM_CLIENTS]\n",
        "federated_train_data = [preprocess(emnist_train.create_tf_dataset_for_client(x))\n",
        "  for x in client_ids\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gNO_Y9j_Lo8X"
      },
      "source": [
        "## Preparing the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LJ0I89ixz8yV"
      },
      "source": [
        "This uses the same model as in the [image classification](federated_learning_for_image_classification.ipynb) tutorial. This model (implemented via `tf.keras`) has a single hidden layer, followed by a softmax layer. In order to use this model in TFF, Keras model is wrapped as a [`tff.learning.models.FunctionalModel`](https://www.tensorflow.org/federated/api_docs/python/tff/learning/models/FunctionalModel). This allows us to perform the model's [forward pass](https://www.tensorflow.org/federated/api_docs/python/tff/learning/models/FunctionalModel#predict_on_batch)\n",
        "aggregator_factory = tff.aggregators.MeanFactory()\n",
        "aggregator = aggregator_factory.create(\n",
        "    model_weights_type.trainable, tff.TensorType(np.float32)\n",
        ")\n",
        "finalizer = tff.learning.templates.build_apply_optimizer_finalizer(\n",
        "    server_optimizer, model_weights_type\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "rw6_YzhedY1f"
      },
      "outputs": [],
      "source": [
        "initializer = tf.keras.initializers.GlorotNormal(seed=0)\n",
        "keras_model = tf.keras.models.Sequential([\n",
        "    tf.keras.layers.Input(shape=(784,)),\n",
        "    tf.keras.layers.Dense(10, kernel_initializer=initializer),\n",
        "    tf.keras.layers.Softmax(),\n",
        "])\n",
        "\n",
        "tff_model = tff.learning.models.functional_model_from_keras(\n",
        "    keras_model,\n",
        "    loss_fn=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
        "    input_spec=federated_train_data[0].element_spec,\n",
        "    metrics_constructor=collections.OrderedDict(\n",
        "        accuracy=tf.keras.metrics.SparseCategoricalAccuracy\n",
        "    ),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LT4ORiiIddoU"
      },
      "source": [
        "## Preparing the optimizers\n",
        "\n",
        "Just as in [`tff.learning.algorithms.build_weighted_fed_avg`](https://www.tensorflow.org/federated/api_docs/python/tff/learning/algorithms/build_weighted_fed_avg), there are two optimizers here: A client optimizer, and a server optimizer. For simplicity, the optimizers will be SGD with different learning rates."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "S6lP8k0Mdf5F"
      },
      "outputs": [],
      "source": [
        "client_optimizer = tff.learning.optimizers.build_sgdm(learning_rate=0.01)\n",
        "server_optimizer = tff.learning.optimizers.build_sgdm(learning_rate=1.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IO5WvBIudhj_"
      },
      "source": [
        "## Defining the building blocks\n",
        "\n",
        "Now that the client work building block, data, model, and optimizers are set up, it remains to create building blocks for the distributor, the aggregator, and the finalizer. This can be done just by borrowing some defaults available in TFF and that are used by FedAvg."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "MKxDA6-qdjHi"
      },
      "outputs": [],
      "source": [
        "@tff.tensorflow.computation\n",
        "def initial_model_weights_fn():\n",
        "  return tff.learning.models.ModelWeights(*tff_model.initial_weights)\n",
        "\n",
        "\n",
        "model_weights_type = initial_model_weights_fn.type_signature.result\n",
        "\n",
        "distributor = tff.learning.templates.build_broadcast_process(model_weights_type)\n",
        "client_work = build_gradient_clipping_client_work(tff_model, client_optimizer)\n",
        "\n",
        "# TFF aggregators use a factory pattern, which create an aggregator\n",
        "# based on the output type of the client work. This also uses a float (the number\n",
        "# of examples) to govern the weight in the average being computed.)\n",
        "aggregator_factory = tff.aggregators.MeanFactory()\n",
        "aggregator = aggregator_factory.create(\n",
        "    model_weights_type.trainable, tff.TensorType(np.float32)\n",
        ")\n",
        "finalizer = tff.learning.templates.build_apply_optimizer_finalizer(\n",
        "    server_optimizer, model_weights_type\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AEYYNHqI1Jif"
      },
      "source": [
        "## Composing the building blocks\n",
        "\n",
        "Finally, you can use a built-in **composer** in TFF for putting the building blocks together. This one is a relatively simple composer, which takes the 4 building blocks above and wires their types together."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "z_86iNeM0IBm"
      },
      "outputs": [],
      "source": [
        "fed_avg_with_clipping = tff.learning.templates.compose_learning_process(\n",
        "    initial_model_weights_fn,\n",
        "    distributor,\n",
        "    client_work,\n",
        "    aggregator,\n",
        "    finalizer\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gcK69pCG16-E"
      },
      "source": [
        "# Running the algorithm\n",
        "\n",
        "Now that the algorithm is done, let's run it. First, **initialize** the algorithm. The **state** of this algorithm has a component for each building block, along with one for the *global model weights*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "Jg22oFx11YKK"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "()"
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "state = fed_avg_with_clipping.initialize()\n",
        "\n",
        "state.client_work"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qmCiEdoq2doJ"
      },
      "source": [
        "As expected, the client work has an empty state (remember the client work code above!). However, other building blocks may have non-empty state. For example, the finalizer keeps track of how many iterations have occurred. Since `next` has not been run yet, it has a state of `0`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "kEuB-8Z71-bd"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "OrderedDict([('learning_rate', 1.0)])"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "state.finalizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2N9XObhZ2zSQ"
      },
      "source": [
        "Now run a training round."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "tKhPuBgW1-3c"
      },
      "outputs": [],
      "source": [
        "learning_process_output = fed_avg_with_clipping.next(state, federated_train_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J7L0jKEe29bk"
      },
      "source": [
        "The output of this (`tff.learning.templates.LearningProcessOutput`) has both a `.state` and `.metrics` output. Let's look at both."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "AMsBmmQz28AZ"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "OrderedDict([('learning_rate', 1.0)])"
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "learning_process_output.state.finalizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hwcfhAbP3VkH"
      },
      "source": [
        "Clearly, the finalizer state has incremented by one, as one round of `.next` has been run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "0K91G_Ob3E05"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "OrderedDict([('distributor', ()), ('client_work', ()), ('aggregator', OrderedDict([('mean_value', ()), ('mean_weight', ())])), ('finalizer', OrderedDict([('update_non_finite', 0)]))])"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "learning_process_output.metrics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2sDyO9uz3Jaz"
      },
      "source": [
        "While the metrics are empty, for more complex and practical algorithms they'll generally be full of useful information."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tPpxe7Ie3gLJ"
      },
      "source": [
        "# Conclusion"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F8uEZw-T3iBB"
      },
      "source": [
        "By using the building block/composers framework above, you can create entirely new learning algorithms, without having to re-do everything from scratch. However, this is only the starting point. This framework makes it much easier to express algorithms as simple modifications of FedAvg. For more algorithms, see [`tff.learning.algorithms`](https://www.tensorflow.org/federated/api_docs/python/tff/learning/algorithms), which contains algorithms such as [FedProx](https://www.tensorflow.org/federated/api_docs/python/tff/learning/algorithms/build_weighted_fed_prox) and [FedAvg with client learning rate scheduling](https://www.tensorflow.org/federated/api_docs/python/tff/learning/algorithms/build_weighted_fed_avg_with_optimizer_schedule). These APIs can even aid implementations of entirely new algorithms, such as [federated k-means clustering](https://www.tensorflow.org/federated/api_docs/python/tff/learning/algorithms/build_fed_kmeans)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "composing_learning_algorithms.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
